#include "TGGmSSL.h"
#include <openssl/sm2.h>
#include <openssl/pem.h>
#include <iostream>

using namespace std;

pair<string, string> 
sm2_genKey() 
{
    EC_KEY *keypair = NULL;
    EC_GROUP *group1 = NULL;

    keypair = EC_KEY_new();
    if(!keypair) {
        cout << "Failed to Gen Key" << endl;
        exit(1);
    }

    group1 = EC_GROUP_new_by_curve_name(NID_sm2p256v1);

    if(group1 == NULL){
        cout << "Failed to Gen Key" << endl;
        exit(1);
    }

    int ret1 = EC_KEY_set_group(keypair, group1);
    if(ret1 != 1){
        cout << "Failed to Gen Key" << endl;
        exit(1);
    }

    int ret2 = EC_KEY_generate_key(keypair);
    if(ret2 != 1){
        cout << "Failed to Gen Key" << endl;
        exit(1);
    }

    size_t pri_len;
    size_t pub_len;
    char *pri_key = NULL;
    char *pub_key = NULL;
 
	BIO *pri = BIO_new(BIO_s_mem());
	BIO *pub = BIO_new(BIO_s_mem());
 
	PEM_write_bio_ECPrivateKey(pri, keypair, NULL, NULL, 0, NULL, NULL);
	PEM_write_bio_EC_PUBKEY(pub, keypair);
 
	pri_len = BIO_pending(pri);
	pub_len = BIO_pending(pub);
 
	pri_key = new char[pri_len + 1];
	pub_key = new char[pub_len + 1];
 
	BIO_read(pri, pri_key, pri_len);
	BIO_read(pub, pub_key, pub_len);
 
	pri_key[pri_len] = '\0';
	pub_key[pub_len] = '\0';
 
	string public_key = pub_key;
	string private_key = pri_key;

	EC_KEY_free(keypair);
	BIO_free_all(pub);
	BIO_free_all(pri);
	delete [] pri_key;
	delete [] pub_key;
 
	return std::pair<string, string>(public_key, private_key);
}

EC_KEY* 
CreateEC(unsigned char* key, int is_public)
{
    EVP_PKEY *pkey = NULL;
	EC_KEY *ec_key = NULL;
	BIO *keybio = NULL;
	keybio = BIO_new_mem_buf(key, -1);
 
	if (keybio==NULL) {
        cout << "Failed to Get Key" << endl;
		return NULL;
	}
 
	if(is_public) {
		pkey = PEM_read_bio_PUBKEY(keybio, NULL, NULL, NULL);
	}
	else {
		pkey = PEM_read_bio_PrivateKey(keybio, NULL, NULL, NULL);
	}

	if(pkey == NULL) {
		return NULL;
	}

    if (!(ec_key = EVP_PKEY_get0_EC_KEY(pkey))
		|| !EC_KEY_is_sm2p256v1(ec_key)) {
		cout << "Failed to Get Key1" << endl;
		return NULL;
	}
 
	if(ec_key == NULL) {
        cout << "Failed to Get Key2" << endl;
		return NULL;
	}
 
	return ec_key;
}

int bio_to_mem(unsigned char **out, int maxlen, BIO *in)
{
    BIO *mem;
    int len, ret;
    unsigned char tbuf[1024];

    mem = BIO_new(BIO_s_mem());
    if (mem == NULL)
        return -1;
    for (;;) {
        if ((maxlen != -1) && maxlen < 1024)
            len = maxlen;
        else
            len = 1024;
        len = BIO_read(in, tbuf, len);
        if (len < -1) {
            fprintf(stderr, "bio error,maxlen:%d,len:%d\n", maxlen, len);
            BIO_free(mem);
            return -1;
        }
        if (len == 0 || len == -1)
            break;
        if (BIO_write(mem, tbuf, len) != len) {
            BIO_free(mem);
            return -1;
        }
        maxlen -= len;

        if (maxlen == 0)
            break;
    }
    ret = BIO_get_mem_data(mem, (char **)out);
    BIO_set_flags(mem, BIO_FLAGS_MEM_RDONLY);
    BIO_free(mem);
    return ret;
}

char *
base64_encode(const char* input, int length, bool with_new_line)
{
	 BIO* bmem = NULL;
	 BIO* b64 = NULL;
	 BUF_MEM* bptr = NULL;
	 b64 = BIO_new(BIO_f_base64());
	  if (!with_new_line)
	 {
	  BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
	 }
	 bmem = BIO_new(BIO_s_mem());
	 b64 = BIO_push(b64, bmem);
	 BIO_write(b64, input, length);
	 BIO_flush(b64);
	 BIO_get_mem_ptr(b64, &bptr);
	 char * buff = (char *)OPENSSL_malloc(bptr->length + 1);
	 memcpy(buff, bptr->data, bptr->length);
	 buff[bptr->length] = 0;
	 BIO_free_all(b64);
	 return buff;
}

int
base64_decode(const char* input, int length, unsigned char* decode, bool with_new_line)
{
	 BIO * b64 = NULL;
	 BIO * bmem = NULL;
	 //char * decode = (char *)malloc(length);
	 memset(decode, 0, length);
	 b64 = BIO_new(BIO_f_base64());
	  if (!with_new_line)
	 {
	  BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
	 }
	 bmem = BIO_new_mem_buf(input, length);
	 bmem = BIO_push(b64, bmem);
	 int len = BIO_read(bmem, decode, length);
     cout << "decode:" << len << endl;
	 BIO_free_all(bmem);
	 return len;
}

int sm2utl_sign(const EVP_MD *md, BIO *in, BIO *out, const char *id,
	ENGINE *e, EC_KEY *ec_key, int sign, int* ret_len)
{
	int ret = 0;
	EVP_MD_CTX *md_ctx = NULL;
	ECDSA_SIG *sig = NULL;
	unsigned char buf[1024];
	size_t siz = sizeof(buf);
	unsigned int ulen = sizeof(buf);
	int len;

	if (!(md_ctx = EVP_MD_CTX_new())
		|| !EVP_DigestInit_ex(md_ctx, md, e)
		|| !SM2_compute_id_digest(md, id, strlen(id), buf, &siz, ec_key)
		|| !EVP_DigestUpdate(md_ctx, buf, siz)) {
		cout << "Error:" << "EVP_MD_CTX_new" << endl;
		goto end;
	}
	while ((len = BIO_read(in, buf, sizeof(buf))) > 0) {
		if (!EVP_DigestUpdate(md_ctx, buf, len)) {
			cout << "Error:" << "BIO_read" << endl;
			goto end;
		}
	}
	if (!EVP_DigestFinal_ex(md_ctx, buf, &ulen)) {
		cout << "Error:" << "EVP_DigestFinal_ex" << endl;
		goto end;
	}
	len = (int)ulen;

	if (sign) {
		unsigned char *p = buf;
		if (!(sig = SM2_do_sign(buf, len, ec_key))
			|| (len = i2d_ECDSA_SIG(sig, &p)) <= 0) {
			cout << "Error:" << "SM2_do_sign" << endl;
			goto end;
		}
	}
	if (BIO_write(out, buf, len) != len) {
		cout << "Error:" << "BIO_write" << endl;
		goto end;
	}

    if(ret_len) {
        *ret_len = len;
    }

	ret = 1;
end:
	EVP_MD_CTX_free(md_ctx);
	ECDSA_SIG_free(sig);
	return ret;
}

int sm2utl_verify(const EVP_MD *md, BIO *in, BIO *out, BIO * sig,
	const char *id, ENGINE *e, EC_KEY *ec_key)
{
	int ret = 0;
	EVP_MD_CTX *md_ctx = NULL;
	unsigned char *sigbuf = NULL;
	unsigned char buf[1024];
	size_t siz = sizeof(buf);
	unsigned int ulen = sizeof(buf);
	int siglen, len;

    siglen = bio_to_mem(&sigbuf, 256, sig);
	if (siglen < 0) {
        cout << "Error:" << "Error reading signature data" << endl;
		goto end;
	}

	if (!(md_ctx = EVP_MD_CTX_new())
		|| !EVP_DigestInit_ex(md_ctx, md, e)
		|| !SM2_compute_id_digest(md, id, strlen(id), buf, &siz, ec_key)
		|| !EVP_DigestUpdate(md_ctx, buf, siz)) {
		cout << "Error:" << "EVP_MD_CTX_new" << endl;
		goto end;
	}
	while ((len = BIO_read(in, buf, sizeof(buf))) > 0) {
		if (!EVP_DigestUpdate(md_ctx, buf, len)) {
			cout << "Error:" << "BIO_read" << endl;
			goto end;
		}
	}
	siz = sizeof(buf);
	if (!EVP_DigestFinal_ex(md_ctx, buf, &ulen)) {
		cout << "Error:" << "EVP_DigestFinal_ex" << endl;
		goto end;
	}
	/* SM2_verify() can check no suffix on signature */
	ret = SM2_verify(NID_undef, buf, ulen, sigbuf, siglen, ec_key);
	if (ret == 1) {
		BIO_puts(out, "Signature Verification Successful\n");
	} else {
        cout << "verify failed:" << ret << endl;
		BIO_puts(out, "Signature Verification Failure\n");
		ret = 0;
	}

end:
	EVP_MD_CTX_free(md_ctx);
	return ret;
}

string 
sm2_sign(const string& private_key, const string& plain_text, const string& id)
{
    const EVP_MD *md = EVP_sm3();
    ENGINE *e = NULL;
    char buf[1024] = {0};
    //char *sign_out = NULL;
    //BIO *b64;
    int ret_len = 0;

    string str_ret;

    BIO* in = BIO_new(BIO_s_mem());
    BIO* out = BIO_new(BIO_s_mem());
    // b64 = BIO_new(BIO_f_base64());
    // out = BIO_push(b64, out);

    EC_KEY *ec_key = CreateEC((unsigned char*)private_key.c_str(), 0);
	if(ec_key == NULL) {
		return string("");
	}

    BIO_write(in, plain_text.c_str(), plain_text.size());

    int ret = sm2utl_sign(md, in, out, id.c_str(), e, ec_key, 1, &ret_len);

    if(ret == 1) {
        // int len = BIO_read(out, buf, 1024);
        // buf[len] = '\0';
        // sign_out = OPENSSL_buf2hexstr(buf, len);
        // cout << "sign out:" << len << "|" << sign_out << endl;
        // str_ret = string(sign_out);
        // OPENSSL_free(sign_out);
        //out = BIO_pop(b64);
        int len = BIO_read(out, buf, ret_len);
        buf[len] = '\0';
        char* b64_str = base64_encode(buf, len, false);
        str_ret = b64_str;
        OPENSSL_free(b64_str);
    }

    //cout << "plain:" << plain_text << "|size:" << plain_text.size() << "|ret_len:" << ret_len << "|ret:" << ret << endl;

    BIO_free_all(in);
	BIO_free_all(out);

    return str_ret;
}

int
sm2_verify(const string& public_key, const string& origin, const string& sig_str, const string& id) 
{
    const EVP_MD *md = EVP_sm3();
    ENGINE *e = NULL;
    unsigned char *decode = (unsigned char *)OPENSSL_malloc(sig_str.length() * sizeof(char));

    BIO* in = BIO_new(BIO_s_mem());
    BIO* out = BIO_new(BIO_s_mem());
    BIO* sig = BIO_new(BIO_s_mem());
    //BIO* b64 = BIO_new(BIO_f_base64());
    BIO_write(in, origin.c_str(), origin.size());

    int len = base64_decode(sig_str.c_str(), sig_str.size(), decode, false);
    int ret = BIO_write(sig, decode, len);

    OPENSSL_free(decode);

    //sig = BIO_push(b64, sig);

    //cout << "sig:" << sig_str << "len:" << sig_str.size() << "|" << len << "|"<< ret << endl;

    EC_KEY *ec_key = CreateEC((unsigned char*)public_key.c_str(), 1);
	if(ec_key == NULL) {
		return 0;
	}

    ret = sm2utl_verify(md, in, out, sig, id.c_str(), e, ec_key);

    BIO_free_all(in);
	BIO_free_all(out);
    BIO_free_all(sig);

    return ret;
}
